// "Copyright [2021] <Copyright Shengkai Lin>"

package src;

import cluster.optical_devices.*;
import cluster.network_devices.*;
import ib_model.src.*;

module Machine 
{
    parameters:
        @controlplane_task;
        int numGpus = default(8);
        int numNvlink_port = default(6);
        int my_index;
        int nic_type = default(0);  // udp-roce = 0; ib = 1
        bool need_wait_topo_machine = default(false);

    gates:
        inout out_port[numGpus];  // Data plane
        inout controller_port;  // Control plane
        inout virtual_data_switch_port[numGpus]; // To send and receive msg from vir data switch

    submodules:
        cpu[numGpus+1] : Cpu 
        {
            parameters:
                my_index = index;
                my_global_index = index + my_index*numGpus;
                need_wait_topo = need_wait_topo_machine;
        }
        
        H[numGpus] : HCA // replace previous nic with HCA
        {
            parameters:
                @display("i=device/server2;p=80,109");
        }
        
        gpu[numGpus] : Gpu 
        {
            parameters:
                numNvlink_port = numNvlink_port;
                my_global_index = index + my_index*numGpus;
                my_index = index;
        }
        

    connections allowunconnected:
        // Only connect cpu0 to the controller directly.
        cpu[0].controller_port$o --> controller_port$o;
        cpu[0].controller_port$i <-- controller_port$i;

        for i=1..(numGpus) 
        {
            cpu[i].controller_port <-->Pcie<--> cpu[0].cpu_port++;
            gpu[i-1].manager_port <-->Pcie<--> cpu[i].gpu_port;
            gpu[i-1].socket_port <-->Pcie<--> H[i-1].gpu_port;
            H[i-1].virtual_data <--> virtual_data_switch_port[i-1];
            H[i-1].port <--> out_port[i-1];
        }

        // The connections of NVlink
        // Connect the basic pattern of nvlink
        // only 1 connection between 2 peers
        gpu[0].nvlink_port[1] <-->Nvlink<--> gpu[1].nvlink_port[3] if numNvlink_port>=4;
        gpu[1].nvlink_port[1] <-->Nvlink<--> gpu[2].nvlink_port[3] if numNvlink_port>=4;
        gpu[2].nvlink_port[1] <-->Nvlink<--> gpu[3].nvlink_port[3] if numNvlink_port>=4;
        gpu[3].nvlink_port[1] <-->Nvlink<--> gpu[0].nvlink_port[3] if numNvlink_port>=4;

        gpu[0].nvlink_port[2] <-->Nvlink<--> gpu[2].nvlink_port[2] if numNvlink_port>=4;
        gpu[1].nvlink_port[2] <-->Nvlink<--> gpu[3].nvlink_port[2] if numNvlink_port>=4;

        gpu[4].nvlink_port[1] <-->Nvlink<--> gpu[5].nvlink_port[3] if numNvlink_port>=4;
        gpu[5].nvlink_port[1] <-->Nvlink<--> gpu[6].nvlink_port[3] if numNvlink_port>=4;
        gpu[6].nvlink_port[1] <-->Nvlink<--> gpu[7].nvlink_port[3] if numNvlink_port>=4;
        gpu[7].nvlink_port[1] <-->Nvlink<--> gpu[4].nvlink_port[3] if numNvlink_port>=4;

        gpu[4].nvlink_port[2] <-->Nvlink<--> gpu[6].nvlink_port[2] if numNvlink_port>=4;
        gpu[5].nvlink_port[2] <-->Nvlink<--> gpu[7].nvlink_port[2] if numNvlink_port>=4;

        for i=0 .. 3 {
            gpu[i].nvlink_port[0] <-->Nvlink<--> gpu[i+4].nvlink_port[0] if numNvlink_port>=4;
        }

        // Connect the extra nvlink port for ring.
        gpu[0].nvlink_port[4] <-->Nvlink<--> gpu[4].nvlink_port[5] if numNvlink_port == 6;
        gpu[4].nvlink_port[4] <-->Nvlink<--> gpu[7].nvlink_port[5] if numNvlink_port == 6;
        gpu[7].nvlink_port[4] <-->Nvlink<--> gpu[6].nvlink_port[5] if numNvlink_port == 6;
        gpu[6].nvlink_port[4] <-->Nvlink<--> gpu[5].nvlink_port[5] if numNvlink_port == 6;
        gpu[5].nvlink_port[4] <-->Nvlink<--> gpu[1].nvlink_port[5] if numNvlink_port == 6;
        gpu[1].nvlink_port[4] <-->Nvlink<--> gpu[2].nvlink_port[5] if numNvlink_port == 6;
        gpu[2].nvlink_port[4] <-->Nvlink<--> gpu[3].nvlink_port[5] if numNvlink_port == 6;
        gpu[3].nvlink_port[4] <-->Nvlink<--> gpu[0].nvlink_port[5] if numNvlink_port == 6;}